#ifndef RBTREE
#define RBTREE

#include <iostream>
#include <utility>

enum class Color
{
    RED,
    BLACK
};

template <typename K, typename V>
struct RBTreeNode
{
    std::pair<K, V> _kv;
    RBTreeNode<K, V> *_left;
    RBTreeNode<K, V> *_right;
    RBTreeNode<K, V> *_parent;
    Color _color;

    RBTreeNode(const std::pair<K, V> &kv)
        : _kv(kv),
          _left(nullptr),
          _right(nullptr),
          _parent(nullptr)
    {
    }
};
/**
 * 红黑树的规则
 * 1. 每个结点不是红⾊就是⿊⾊
 * 2. 根结点是⿊⾊的
 * 3. 如果⼀个结点是红⾊的，则它的两个孩⼦结点必须是⿊⾊的，也就是说任意⼀条路径不会有连续的红⾊结点。
 * 4. 对于任意⼀个结点，从该结点到其所有NULL结点的简单路径上，均包含相同数量的⿊⾊结点
 */

template <typename K, typename V>
class RBTree
{
    using Node = RBTreeNode<K, V>;
    const Color red = Color::RED;
    const Color black = Color::BLACK;

private:
    void RotateRight(Node *parent)
    {
        /**
         * 假设某一个子树新插入节点后不平衡了，把这个子树的父亲节点旋转到根节点上
         * 细节：子节点的判空，非空才能使用，parent指针的控制
         * 判断这是单独的子树还是根开始的子树，根据parent指针判断
         * 提前记录parent指针
         * 平衡因子的更新，只有子树高度发生变化了才需要更新
         */

        Node *subL = parent->_left;
        Node *subLR = subL->_right;

        Node *parent_ = parent->_parent;

        parent->_left = subLR;
        if (subLR)
            subLR->_parent = parent;

        subL->_right = parent;
        parent->_parent = subL;

        if (parent_ == nullptr)
        {
            _root = subL;
            subL->_parent = nullptr;
        }
        else
        {
            // 判断这是左子树还是右子树
            if (parent_->_left == parent)
            {
                parent_->_left = subL;
            }
            else
            {
                parent_->_right = subL;
            }
            subL->_parent = parent_;
        }
    }

    void RotateLeft(Node *parent)
    {
        /**
         * 和右单旋逻辑一样
         */

        Node *subR = parent->_right;
        Node *subRL = subR->_left;

        Node *parent_ = parent->_parent;

        parent->_right = subRL;

        if (subRL)
            subRL->_parent = parent;

        subR->_left = parent;
        parent->_parent = subR;

        if (parent_ == nullptr)
        {
            _root = subR;
            subR->_parent = nullptr;
        }
        else
        {
            if (parent == parent_->_left)
            {
                parent_->_left = subR;
            }
            else
            {
                parent_->_right = subR;
            }
            subR->_parent = parent_;
        }
    }

    bool check(Node *root, int curblacknum, const int &blacknum)
    {
        // 递归检查,一条路径走完了
        if (root == nullptr)
        {
            if (curblacknum != blacknum)
            {
                std::cout << "这不是一个红黑树，不符合任意路径上的黑色节点数量相同" << std::endl;
                return false;
            }
            return true;
        }

        // 一个节点有两个孩子，不好检查，检查一个孩子及其父亲是否是连续的红色即可
        if (root->_color == red && root->_parent->_color == red)
        {
            std::cout << "存在连续红色节点" << std::endl;
            return false;
        }

        if (root->_color == black)
            curblacknum++;
        std::cout << "hello" << std::endl;

        return check(root->_left, curblacknum, blacknum) && check(root->_right, curblacknum, blacknum);
    }

    void _Inorder(Node *root)
    {
        if (root == nullptr)
            return;
        _Inorder(root->_left);
        std::cout << root->_kv.first << ":" << root->_kv.second << std::endl;
        _Inorder(root->_right);
    }

public:
    RBTree() {}
    ~RBTree() {}

    void Inorder()
    {
        _Inorder(_root);
    }

    bool insert(const std::pair<K, V> &kv)
    {
        // 空树插入
        if (_root == nullptr)
        {
            _root = new Node(kv);
            _root->_color = black;
            return true;
        }

        // 非空树插入
        Node *cur = _root;
        Node *parent_ = nullptr;
        // 先找到可以插入的位置，本质还是一个搜索树，还是符合左节点值小于根节点，右节点值大于根节点
        while (cur)
        {
            if (cur->_kv.first > kv.first)
            {
                parent_ = cur;
                cur = cur->_left;
            }
            else if (cur->_kv.first < kv.first)
            {
                parent_ = cur;
                cur = cur->_right;
            }
            else
                return false;
        }

        // 判断父节点的颜色
        // 需要保证规则4不被破坏,规则4很难维护，我们尽量保证不要新增黑色节点
        // if (cur->_parent->_color == red)
        //     cur->_color = black;
        // else
        // 强制为红色
        cur = new Node(kv);
        cur->_color = red;

        // 判断这个节点放在哪里
        if (parent_->_kv.first < kv.first)
        {
            parent_->_right = cur;
        }
        else
        {
            parent_->_left = cur;
        }

        cur->_parent = parent_;
        // 向上调整节点颜色,关键看叔叔，新插入的节点必须是红色，才能保证不破坏规则4
        /**
         * 变色
         */
        // 获取爷叔节点

        while (parent_ && parent_->_color == red)
        {
            Node *grandfather_ = parent_->_parent;

            if (parent_ == grandfather_->_left)
            {
                Node *uncle = grandfather_->_right;

                if (uncle && uncle->_color == red)
                {
                    parent_->_color = uncle->_color = black;
                    grandfather_->_color = red;

                    // 继续向上处理
                    cur = grandfather_;
                    parent_ = cur->_parent;
                }

                // 叔叔节点为空，或者为黑色，需要旋转加变色
                else
                {
                    // 是左边节点，就是单纯的左边高，需要优选
                    if (cur == parent_->_left)
                    {
                        RotateRight(grandfather_);
                        parent_->_color = black;
                        grandfather_->_color = red;
                    }
                    else
                    {
                        // 右边节点，需要旋转两次加变色
                        RotateLeft(parent_);
                        RotateRight(grandfather_);

                        cur->_color = black;
                        grandfather_->_color = red;
                    }
                    break;
                }
            }
            else
            {
                Node *uncle = grandfather_->_left;
                if (uncle && uncle->_color == red)
                {
                    parent_->_color = uncle->_color = black;
                    grandfather_->_color = red;

                    cur = grandfather_;
                    parent_ = cur->_parent;
                }
                else
                {
                    // 这是在右子树，需要注意
                    if (cur == parent_->_right)
                    {
                        // 单纯右边高就左旋
                        RotateLeft(grandfather_);
                        parent_->_color = black;
                        grandfather_->_color = red;
                    }
                    else
                    {
                        RotateRight(parent_);
                        RotateLeft(grandfather_);

                        cur->_color = black;
                        grandfather_->_color = red;
                    }
                    break;
                }
            }
        }
        _root->_color = black;
        return true;
    }

    Node *find(const K &key)
    {
        Node *cur = _root;
        while (cur)
        {
            if (cur->_kv.first > key)
            {
                cur = cur->_left;
            }
            else if (cur->_kv.first < key)
            {
                cur = cur->_right;
            }
            else
                return cur;
        }
        return nullptr;
    }

    bool IsBalance(Node *root)
    {
        // 按照规则来写
        if (root == nullptr)
            return true;

        if (root->_color != black)
            return false;

        // 随机选一个节点进行路径上黑色节点的计算，然后比较其余的节点路径上的黑色节点的数量
        Node *temp = root;
        int blacknum = 0;
        while (temp)
        {
            if (temp->_color == black)
                blacknum++;
            temp = temp->_left;
        }
        // std::cout << "hello" << std::endl;
        return check(root, 0, blacknum);
    }

    Node *getRoot() { return _root; }

private:
    Node *_root = nullptr;
};

#endif